#include <iostream>
#include <vector>
#include "lib/binary_tree.h"

BinaryTreeWithParentNode* FindNode(BinaryTreeWithParentNode* cur_node) {
  if (cur_node == nullptr) {
    return nullptr;
  }

  if (cur_node->right) {
    auto next_node = cur_node->right;
    while (next_node->left) {
      next_node = next_node->left;
    }
    return next_node;
  } else if (cur_node->parent) {
    auto parent_node = cur_node->parent;
    while (parent_node && cur_node == parent_node->right) {
      cur_node = cur_node->parent;
      parent_node = parent_node->parent;
    }
    return parent_node;
  }
  return nullptr;
}

int main() {
  auto node1 = CreateBinaryTreeWithParentNode(1);
  auto node2 = CreateBinaryTreeWithParentNode(2);
  auto node3 = CreateBinaryTreeWithParentNode(3);
  auto node4 = CreateBinaryTreeWithParentNode(4);
  auto node5 = CreateBinaryTreeWithParentNode(5);
  auto node6 = CreateBinaryTreeWithParentNode(6);
  auto node7 = CreateBinaryTreeWithParentNode(7);
  auto node8 = CreateBinaryTreeWithParentNode(8);
  auto node9 = CreateBinaryTreeWithParentNode(9);

  ConnectNode(node1, node2, node3);
  ConnectNode(node2, node4, node5);
  ConnectNode(node3, node6, node7);
  ConnectNode(node6, node8, node9);

  PrintBinaryTree(node1);

  auto find_node2 = FindNode(node4);
  PrintBinaryTreeNode(find_node2);

  auto find_node8 = FindNode(node1);
  PrintBinaryTreeNode(find_node8);

  auto find_node3 = FindNode(node9);
  PrintBinaryTreeNode(find_node3);

  return 0;
}
